import numpy as np
import cv2
import os
import os.path as osp
import json
import re
import copy
import shutil
from tqdm import tqdm
from multiprocessing import Pool
from PIL import Image, ImageColor, ImageFont, ImageDraw, ImageFilter

from tool import filesystem

# from tool import darknet_tool # export PYTHONPATH=$PYTHONPATH:`pwd`

def split_polygon_to_via(data_root,resize_scale, split_w,ori_via_name, dst_via_name):
    """
        将长文本切割成小窄框,
        供训练yolo使用,
        并保存到同级目录下,
    """
    via_files = filesystem.get_all_filepath(data_root, [ori_via_name])
    thread_count  = 12
    # print('Parent process %s.' % os.getpid())
    p = Pool(thread_count)
    for f in via_files:
        path_dir = os.path.dirname(f)
        if not os.path.isdir(path_dir):
            continue
        print(f)
        p.apply_async(__split_polygon_to_via_impl, args=(path_dir, resize_scale,split_w,ori_via_name, dst_via_name))
    # print('Waiting for all subprocesses done...')
    p.close()
    p.join()


def deal_many_dir(data_root, save_dir,via_name='', save_image=False,gen_type=""):
    """
        制作darknet训练文件,
        将via转换成txt
    """
    via_files = filesystem.get_all_filepath(data_root, [via_name])
    thread_count  = 8
    # print('Parent process %s.' % os.getpid())
    # p = Pool(thread_count)
    for f in via_files:
        path_dir = os.path.dirname(f)
        if not os.path.isdir(path_dir):
            continue
        print(f)
        __convert_via_to_darknet(f, data_root, save_dir,via_name,save_image, gen_type)
        # p.apply_async(__convert_via_to_darknet, args=(path_dir, save_dir,via_name,save_image, gen_type))
    print('Waiting for all subprocesses done...')
    # p.close()
    # p.join()

def deal_many_dir_local(data_root,via_name='', gen_type=""):
    """
        就地处理
        制作darknet训练文件,
        将via转换成txt
    """
    via_files = filesystem.get_all_filepath(data_root + os.sep+"images", [via_name])

    save_dir = data_root

    thread_count  = min(12, len(via_files)) 
    # print('Parent process %s.' % os.getpid())
    p = Pool(thread_count)
    for f in via_files:
        path_dir = os.path.dirname(f)
        if not os.path.isdir(path_dir):
            continue
        print(f)
        # __convert_via_to_darknet(f, data_root, save_dir,via_name,False, gen_type)

        p.apply_async(__convert_via_to_darknet, args=(f, data_root, save_dir,via_name,False, gen_type))
    # print('Waiting for all subprocesses done...')
    p.close()
    p.join()


def check_darket_train_data(save_dir, num=10, ext=".jpg"):
    """
        可视化darknet训练数据
    """
    val_files = []
    with open(save_dir+os.sep + "val.txt", "r") as rf:
        for l in rf.readlines():
            val_files.append(l.strip())
        
    np.random.shuffle(val_files)
    np.random.shuffle(val_files)

    for image_path in val_files[:num]:
        image = cv2.imread(image_path)
        print("image_path: ", image_path)
        height, width = image.shape[:2]

        label_path = image_path.replace("{}{}images{}".format(save_dir, osp.sep, osp.sep), "{}{}labels{}".format(save_dir, osp.sep, osp.sep)).replace(ext, ".txt")
        with open(label_path, "r") as rf:
            for l in rf.readlines():
                if l.replace("\n", "") == "":break

                infos = l.replace("\n", "").split(" ")
                w, h = float(infos[3]) * width, float(infos[4]) * height
                x1 = float(infos[1]) * width - w/2
                y1 = float(infos[2]) * height - h/2
                image = cv2.rectangle(image, (int(x1),int(y1)), (int(x1+w), int(y1+h)), (255, 0,0))
        cv2.imwrite(save_dir+os.sep+os.path.basename(image_path), image)


def create_train_val_txt(data_root, yolov8=True):
    """
    制作训练测试集
        use .txt to random and create dataset
        @data_root: coco格式的数据集存放目录
    """
    all_files = filesystem.get_all_filepath(data_root + os.sep + "labels", [".txt"])
    length = len(all_files)

    indeces = list(np.arange(length))
    np.random.shuffle(indeces)
    np.random.shuffle(indeces)

    train_txt = open(data_root + os.sep +"train.txt", "w")
    val_txt = open(data_root + os.sep +"val.txt", "w")

    for i in range(length):
        img_path = all_files[indeces[i]].replace("{}labels{}".format(osp.sep, osp.sep), "{}images{}".format(osp.sep, osp.sep)).replace(".txt", ".jpg")
        if i < length * 0.05:
            val_txt.write(img_path + "\n")
        else:
            train_txt.write(img_path + "\n")

        # yolov8中背景图片不需要有label_path
        if yolov8:
            background = False
            with open(all_files[indeces[i]], "r") as rf:
                for l in rf.readlines():
                    if len(l) < 2:
                        background = True
                        print("background,", all_files[indeces[i]])
            if background:
                os.remove(all_files[indeces[i]])

def __polylines(img,points):
    im = np.zeros(img.shape[:2], dtype = "uint8")
    for point in points:
        b = np.array([point],dtype = np.int32)
        cv2.fillPoly(im, b, 255)
    return im

def __check_points(points,w,h):
    ##检测标注是否正确
    check=False
    for point in points:
        for x,y in point:
            if x >= w or y >= h or x<0 or y<0:
                check=True
                break
    return check

def __get_points(label_data):
    points =[]
    for line in label_data["regions"]:
        
        points.append([(x, y) for x,y in zip(line['shape_attributes']["all_points_x"],
                        line['shape_attributes']["all_points_y"])])
    return points

def __get_rects(label_data, img_width, img_height, filter_label=False):
    rects =[]
    for line in label_data["regions"]:
        x = line['shape_attributes']["x"]
        y = line['shape_attributes']["y"]
        w = line['shape_attributes']["width"]
        h = line['shape_attributes']["height"]

        if x < 0:
            w = w + x
            x = 0

        if y < 0:
            h = h + y
            y = 0


        rects.append([x,
                        y,
                        img_width - x if x + w > img_width else  w,
                        img_height - y if y + h > img_height else  h,
                        "rect" if filter_label else line["region_attributes"].get("label", None) 
                    ]
                    )
    return rects

def __get_rects_circle_meter(label_data):
    rects =[]
    for line in label_data["regions"]:
        rects.append([line['shape_attributes']["x"],
                        line['shape_attributes']["y"],
                        line['shape_attributes']["width"],
                        line['shape_attributes']["height"],
                        line["region_attributes"]["label"]
                    ]
                    )
    return rects


def __get_label_int(gen_type, label):
    if gen_type == "road_vihicle.merge":
        if label == "car":
            return 0
        elif label == "van":
            return 1
        elif label == "bus":
            return 2
        elif label == "truck":
            return 3
        elif label == "closed-tricycle":
            return 4
        elif label == "open-tricycle":
            return 5
        elif label == "forklift":
            return 6
        elif label == "垃圾箱满溢":
            return 7
        elif label == "防撞桶正常":
            return 8
        elif label == "电力设施异常":
            return 9
        elif label == "沿街晾挂":
            return 10
        elif label == "护栏正常":
            return 11
        elif label == "电力设施正常":
            return 12
        elif label == "无照游商":
            return 13
        elif label == "暴露垃圾":
            return 14
        elif label == "擅自搭建气模、拱门":
            return 15
        elif label == "井盖正常":
            return 16
        elif label == "占道经营":
            return 17
        elif label == "乱堆物堆料":
            return 18
        elif label == "店外经营":
            return 19
        elif label == "违规广告":
            return 20
        elif label == "移动护栏":
            return 21
        elif label == "垃圾桶正常":
            return 22
        elif label == "便道桩异常":
            return 23
        elif label == "便道桩正常":
            return 24
        elif label == "道路积水":
            return 25
        elif label == "消防龙头正常":
            return 26
        elif label == "井盖缺失":
            return 27
        elif label == "护栏异常":
            return 28
        elif label == "道路不洁":
            return 29
        elif label == "道路破损":
            return 30
        elif label == "立杆异常":
            return 31
        elif label == "通讯交接箱正常":
            return 32
        elif label == "道路垃圾":
            return 33
        elif label == "行道树异常":
            return 34
        elif label == "横幅乱吊挂":
            return 35
        elif label == "占道施工":
            return 36
        elif label == "违规宣传标语":
            return 37
        elif label == "共享单车乱停":
            return 38
        elif label == "违章接坡":
            return 39
        elif label == "废弃家具设备":
            return 40
        elif label == "非法小广告":
            return 41
        elif label == "防撞桶异常":
            return 42
        else:
            print("FileExistsError....")
            raise FileExistsError
    if gen_type == "road_vihicle":
        if label == "car":
            return 0
        elif label == "van":
            return 1
        elif label == "bus":
            return 2
        elif label == "truck":
            return 3
        elif label == "closed-tricycle":
            return 4
        elif label == "open-tricycle":
            return 5
        elif label == "forklift":
            return 6
        else:
            print("FileExistsError....")
            raise FileExistsError
    if gen_type == "car":
        if label == "person":
            return 0
        elif label == "motorcycle":
            return 1
        elif label == "car":
            return 2
        elif label == "pickup":
            return 3
        elif label == "bus":
            return 4
        elif label == "truck":
            return 5
        elif label == "big_truck":
            return 6
        else:
            raise FileExistsError
    if gen_type == "tt100k":
        return 0
    if gen_type == "vote":
        return 0
    if gen_type == "falan":
        return 0
    if gen_type == "sewage":
        if label == "sewage":
            return 0
        elif label == "water":
            return 1
        else:
            raise FileExistsError
    if gen_type == "waterlogging":
        return 0
    if gen_type == "dust":
        return 0
    if gen_type == "garbage":
        return 0
    if gen_type == "boat_smoke":
        return 0
    if gen_type == "ringelman":
        if label == "scar":
            return 0
        elif label == "bcar":
            return 0
        elif label == "car":
            return 0
        elif label == "ringelman":
            return 1
        else:
            print("gen_type:{}, label:{}".format(gen_type, label))
            raise FileExistsError
    elif gen_type == "boeing":
        if label == "other":
            return 0
        else :
            return 1
    elif gen_type == "face_mask":
        if label == "face":
            return 0
        elif label == "face_mask":
            return 1
        else :
            raise FileExistsError
    elif gen_type == "beijing_handwriting":
        if label == "pinyin":
            return 0
        elif label == "char":
            return 1
        else:
            raise FileExistsError
    elif gen_type == "ct_meter":
        if label == "meter":
            return 0
        elif label == "switch":
            return 1
        elif label == "water":
            return 2
        else:
            raise FileExistsError
    elif gen_type == "fire":
        if label == "smoke":
            return 0
        elif label == "fire":
            return 1
        else:
            raise FileExistsError
    elif gen_type == "ct_water":
        if label == "water":
            return 0
        else:
            raise FileExistsError
    elif gen_type == "bkc_meter":
        if label == "meter":
            return 0
        elif label == "white":
            return 1
        elif label == "black":
            return 2
        else:
            raise FileExistsError
    elif gen_type == "tyre":
        if label == "tyre":
            return 0
        elif label == "hub":
            return 1

        # if label == "xi":
        #     return 0
        # elif label == "kuan":
        #     return 1
        # elif label == "black":
        #     return 2
        # elif label == "white":
        #     return 3
        else:
            raise FileExistsError
    elif gen_type == "circle_meter":
        if label == "0":
            return 0
        elif label == "1":
            return 1
        elif label == "2":
            return 2
        elif label == "3":
            return 3
        elif label == "4":
            return 4
        elif label == "5":
            return 5 
        elif label == "6":
            return 6
        elif label == "7":
            return 7
        elif label == "8":
            return 8
        elif label == "9":
            return 9
        else:
            print("-------------------", label)
            raise FileExistsError
    else:
        raise FileExistsError

    
def __get_polygon_to_rect(label_data, filter_label=False):
    """
        标签数据是points, 返回points的拟合矩形
    """
    points =[]
    labels = []
    for line in label_data["regions"]:
        
        points.append([[x, y] for x,y in zip(line['shape_attributes']["all_points_x"],
                        line['shape_attributes']["all_points_y"])])
        labels.append("rect" if filter_label else line["region_attributes"].get("label", None))

    rects =[]
    for idx, polygon in enumerate(points):
        rect = list(cv2.boundingRect(np.array(polygon)))
        rect.append(labels[idx])
        rects.append(rect)
    return rects

def __get_polygon_to_rect_for_plate_2(label_data):
    """
        标签数据是points, 返回points的拟合矩形
    """
    points =[]
    labels = []
    for line in label_data["regions"]:

        if line["region_attributes"].get("label", None):
            label = line["region_attributes"].get("label")
            if len(label) == 1:
                continue
            label = "rect"

            labels.append(label)

                    
            points.append([[x, y] for x,y in zip(line['shape_attributes']["all_points_x"],
                            line['shape_attributes']["all_points_y"])])
        

    rects =[]
    for idx, polygon in enumerate(points):
        rect = list(cv2.boundingRect(np.array(polygon)))
        rect.append(labels[idx])
        rects.append(rect)
    return rects



def __letterbox_image(cv_image, std_w, std_h):

    image_h, image_w = cv_image.shape[:2]
    
    # if max(image_w, image_h) < min(size):
    if image_w <= std_w and image_h <= std_h:
        resized_image = cv_image
        new_w = std_w
        new_h = std_h
        fw =1.0
        fh =1.0
    else:
        new_w = int(image_w * min(std_w*1.0/image_w, std_h*1.0/image_h))
        new_h = int(image_h * min(std_w*1.0/image_w, std_h*1.0/image_h))
        resized_image = cv2.resize(cv_image, (new_w,new_h), cv2.INTER_AREA)
        fw = image_w / new_w
        fh = image_h / new_h

    padded_image=cv2.copyMakeBorder(resized_image, 0,std_h-new_h, 0, std_w-new_w, cv2.BORDER_CONSTANT, value=[255,255,255])
    return padded_image, fw, fh

def __resize_im(img, scale=416, max_scale=608):

    img,fw, fh = __letterbox_image(img,scale,scale )
    # h,w = img.shape[:2]
    # f=float(scale)/min(h, w)
    # if max_scale is not None:
    #     if  f*max(h, w)>max_scale:
    #         f=float(max_scale)/max(h, w)
    # newW,newH = int(w*f),int(h*f)
    # newW,newH = newW-(newW%32),newH-(newH%32)
    # fw = w/newW
    # fh = h/newH
    # tmpImg = cv2.resize(img, None, None, fx=1/fw, fy=1/fh, interpolation=cv2.INTER_LINEAR)
    return img,fw,fh

def __cleam_im(im):
        avg = 127
        im[im>avg]=255
        im[im<=avg]=0
        y,x = np.where(im==255)
        xmin,ymin,xmax,ymax = (min(x),min(y),max(x),max(y))
        return xmin,ymin,xmax,ymax
    
def __adjust_height(h):
    """
    调整box高
    """
    heights = [11, 16, 23, 33, 48, 68, 97, 139, 198, 283]
    N = len(heights)
    for i in range(N-1):
        if h<=heights[i]+heights[i]*0.44/2:
            return heights[i]
    return h

def __img_split_to_box(im,label,splitW = 15,adjust=True):
    """
    均等分割box
    """
    tmpIm = im==255
    h,w = tmpIm.shape[:2]
    num = w//splitW+1

    if w % splitW >= splitW / 2:
        num += 1

    mean_height = 0
    box = []
    for i in range(num-1):
        # if i == num -2:
        #     xmin,ymin,xmax,ymax = w- splitW,0,w,h,
        # else:
        xmin,ymin,xmax,ymax = splitW*i,0,splitW*(i+1),h
        ##迭代寻找最优ymin,ymax
        childIm = tmpIm[ymin:ymax,xmin:xmax]
        checkYmin = False
        checkYmax = False
        for j in range(ymax):
            if not checkYmin:
                if childIm[j].max():
                    ymin = j
                    checkYmin = True
            if not checkYmax:
                if childIm[ymax-j-1].max():
                    ymax = ymax-j
                    checkYmax = True
               
        if adjust:
            childH = ymax-ymin+1
            cy = (ymax+ymin)/2
            # childH = adjust_height(childH)
            ymin = cy-childH/2
            ymax = cy+childH/2     
            if 2< i < num-1-2:
                mean_height = (mean_height+ ymax-ymin ) / 2

        box.append([xmin,ymin,xmax,ymax, label])
    
    # 添加最后一个， 以最右边为结束
    # if w % splitW != 0:
    #     box.append([w- splitW,box[-1][1],w,box[-1][3], label])

    # 删除切分小框两边的很低的框
    if box[num-2][3] - box[num-2][1] < mean_height* 0.6:
        del box[num-2]
    if box[num-3][3] - box[num-3][1] < mean_height* 0.6:
        del box[num-3]
    if box[1][3] - box[1][1] < mean_height* 0.6:
        del box[1]
    if box[0][3] - box[0][1] < mean_height* 0.6:
        del box[0]
    return box

def __resize_img_box(dataRoot, label_data,scale=416, max_scale=608,splitW=15,adjust=True):
    path = dataRoot + os.sep + label_data["filename"]
    if not os.path.exists(path):
        print("continue ", path)
        return None, None

    img = cv2.imread(path)

    points = __get_points(label_data)
    if len(points) == 0:
        return None, None
    h,w = img.shape[:2]
    check = __check_points(points,w,h)
    if  check:
        print("check_points None")
        return None,None
    # img,fw,fh = __resize_im(img, scale=scale, max_scale=max_scale)
    fw , fh= 1.0,1.0

    # print("fw: ", fw, " fh: ", fh)
    boxes = []
    # for rect in rects:
    #     new_rect = [int(rect[0]/fw), int(rect[1]/fh),int(rect[2]/fw), int(rect[3]/fh), rect[4]]
    #     im = polylines(img,new_rect)
    for point in points:
        point = [[bx[0]/fw,bx[1]/fh] for bx in point]
        im = __polylines(img,[point])
        if im.max()==0:
            continue
            
        xmin,ymin,xmax,ymax = __cleam_im(im)
        tmp = im[ymin:ymax,xmin:xmax]
        box = __img_split_to_box(tmp,"rect",splitW = splitW,adjust=adjust)
        childBoxes = []
        for bx in box:
            xmin_,ymin_,xmax_,ymax_, label = bx

            xmin_,ymin_,xmax_,ymax_ = xmin+xmin_,ymin_+ymin,xmax_+xmin,ymax_+ymin
            boxes.append([xmin_,ymin_,xmax_,ymax_, label])

    return img,boxes
    
def __convert(size, box):
    """
    box = xmin,ymin,xmax,ymax
    转化为darknet训练所需格式 cx,cy,w,h
    """
    dw = 1./(size[0])
    dh = 1./(size[1])
    x = (box[0] + box[2])/2.0 - 1
    y = (box[1] + box[3])/2.0 - 1
    w = box[2] - box[0]
    h = box[3] - box[1]
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    return [str(x),str(y),str(w),str(h)]


def __convert_annotation(dataRoot, label_data,scale=608, max_scale=1024,splitW=8,adjust=False):
    img,boxes = __resize_img_box(dataRoot,label_data,scale=scale, max_scale=max_scale,splitW=splitW,adjust=adjust)
    if img is None or len(boxes)==0:
        return None,''
    h,w = img.shape[:2]
    regions = []
    for bx in boxes:
        attributes = {}
        attributes["shape_attributes"] ={"name":"rect","x":int(bx[0]),"y":int(bx[1]),"width":int(bx[2]-bx[0]),"height":int(bx[3]-bx[1])}
        # attributes["region_attributes"] ={"label":bx[4]}
        attributes["region_attributes"] ={"label":"rect"}
        regions.append(attributes)
    label_data["regions"] = regions
    return img,label_data

# def write_for_via(img,newBoxes,filename):
#     imgP = os.path.join(JPEGP,filename+'.jpg')
#     txtP = os.path.join(labelP,filename+'.txt')
#     cv2.imwrite(imgP,img)
#     with open(txtP,'w') as f:
#         f.write(newBoxes)
        

def __convert_to_via(data_ori_path, via_path):
    with open(data_ori_path,"r") as rf:
        data_dict = json.loads(rf.read())

    with open(via_path,"r") as rf:
        via_dict = json.loads(rf.read())

    save_dict = dict()
    reg = re.compile(r"(.*\.jpg)(\d+)")
    for key in via_dict.keys():
        outs = reg.findall(key)
        if len(outs) == 0:
            print("con't match re")
            continue
        filename = outs[0][0]
        size = int(outs[0][1])
        if data_dict.get(filename, None) is None:
            continue
        data_dict[filename]["size"] = size
        save_dict[filename + str(size)] = data_dict[filename]

    with open(os.path.dirname(data_ori_path) + os.sep + "via_region_data.json","w") as wf:
        wf.write(json.dumps(save_dict))



def __split_polygon_to_via_impl(data_root, resize_scale,split_w, ori_via_name, dst_via_name):
    region_data = data_root + os.sep +ori_via_name
    with open(region_data) as f:
        train_labels = json.loads(f.read())

    ##创建保存目录    
    # save_root = save_dir
    # os.makedirs(save_root, exist_ok=True)
    save_dict = {}

    for p in train_labels.keys():
        img,newBoxes_json = __convert_annotation(data_root, train_labels[p], scale=resize_scale, max_scale=resize_scale, splitW=split_w, adjust=True)
        if img is None :
            print("continue: ", train_labels[p])
            continue

        save_dict[train_labels[p]["filename"]] = newBoxes_json
        # # 保存图片
        # save_path = save_root + os.sep + train_labels[p]["filename"]
        # cv2.imwrite(save_path, img)
 
    save_json_path = data_root + os.sep +dst_via_name
    with open(save_json_path,'w') as f:
        f.write(json.dumps(save_dict))



def __convert_via_to_darknet(via_file_path, dataRoot, save_dir,via_name='', save_image=True,gen_type=""):
    if not os.path.exists(via_file_path):
        print("return :", via_file_path)
        return
    with open(via_file_path, "r", encoding="utf-8") as f:
        train_labels = json.loads(f.read())
    img_dir = os.path.dirname(via_file_path)

    ##创建保存目录    
    save_image_root = save_dir + os.path.dirname(via_file_path).replace(dataRoot, "")
    os.makedirs(save_image_root, exist_ok=True)

    save_label_root = save_dir +os.sep + "labels" + os.path.dirname(via_file_path).replace(dataRoot + os.sep +"images", "")
    os.makedirs(save_label_root, exist_ok=True)

    for p in train_labels.keys():

        img_path = img_dir+os.sep+train_labels[p]["filename"]
        if not os.path.exists(img_path):
            print("continue img_path ", img_path)
            continue
        cv_image = cv2.imread(img_path)
        h,w = cv_image.shape[:2]


        if gen_type in ["beijing_handwriting"]:
            rects = __get_polygon_to_rect(train_labels[p])
        elif gen_type in [ "common_ocr", "driving_license"]:
            rects = __get_polygon_to_rect(train_labels[p])
        elif gen_type in ["plate_text"]:
            rects = __get_polygon_to_rect(train_labels[p], True)
        elif gen_type in [ "plate_recognition_2"]:
            rects = __get_polygon_to_rect_for_plate_2(train_labels[p])
        elif gen_type in ["dust", "plate_recognition", "white_black_meter"]:
            rects = __get_rects(train_labels[p],w,h, True)
        elif gen_type in [ "tyre"]:
            rects = __get_rects(train_labels[p],w,h)
        elif gen_type in [ "ct_meter", "ct_water"]:
            rects = __get_rects(train_labels[p],w,h)
        elif gen_type in [ "bkc_meter"]:
            rects = __get_rects(train_labels[p],w,h)
        elif gen_type == "circle_meter":
            rects = __get_rects_circle_meter(train_labels[p])
        elif gen_type == "boeing":
            rects = __get_rects(train_labels[p],w,h, filter_label=False)
        elif gen_type in ["road_vihicle.merge", "road_vihicle", "vote", "garbage", "fire", "face_mask", "waterlogging", "ringelman", "sewage", "tt100k", "falan"]:
            rects = __get_rects(train_labels[p],w,h, filter_label=False)
        elif gen_type in ["boat_smoke"]:
            rects = __get_rects(train_labels[p],w,h, filter_label=True)
        elif gen_type in ["car"]:
            rects = __get_rects(train_labels[p],w,h, False)
        else:
            raise NotImplementedError
        if len(rects) == 0 :
            print("continue ", p)
            continue


        # print(h,w)
        with open(save_label_root + os.sep + ".".join(train_labels[p]["filename"].split(".")[:-1])+".txt", "w") as wf:
            if rects[0][4] == "retain":        # 保留背景图片 
                wf.write("\n")
            else:
                # print(img_path)
                for rect in rects:
                    string = "{} {} {} {} {}\n".format(__get_label_int(gen_type, rect[4]), (rect[0] +rect[2]/2) /w, (rect[1]+rect[3]/2)/h, rect[2]/w, rect[3]/h)
                    wf.write(string)
        
        if save_image:
            save_image_path = save_image_root + os.sep + train_labels[p]["filename"]
            shutil.copy(img_dir+os.sep+train_labels[p]["filename"], save_image_path)


def __move_image(data_dir, save_dir):
    for p in os.listdir(data_dir):
        new_dir = data_dir + os.sep + p

        save_path = save_dir + os.sep + p
        os.makedirs(save_path, exist_ok=True)

        for name in os.listdir(new_dir):
            if name.endswith(".jpg"):
                
                save_file = save_path + os.sep +  name

                if np.random.randint(2) == 0:
                    image = Image.open(new_dir+os.sep+name)
                    image = image.resize((512,512), Image.BICUBIC if np.random.randint(2) == 0 else Image.LINEAR)
                    image.save(save_file)
                else:
                    image = cv2.imread(new_dir+os.sep+name)
                    image = cv2.resize(image, (512,512), image, 0., 0., cv2.INTER_CUBIC if np.random.randint(2) == 0 else cv2.INTER_AREA )
                    cv2.imwrite(save_file, image)
        src_json = new_dir + os.sep + "via_region_data_split.json"
        dst_json = save_path + os.sep+ "via_region_data_split.json"
        shutil.move(src_json, dst_json)


def __remove_image(data_dir, extension=".jpg", flag=".txt"):
    for p in os.listdir(data_dir):
        new_dir = data_dir + os.sep + p


        for name in os.listdir(new_dir):
            if name.endswith(extension):
                file_path = new_dir+ os.sep + name.replace(extension, flag)
                file_path = file_path.replace("images", "labels")
                if not os.path.exists(file_path):
                    os.remove(new_dir+ os.sep + name)
                    print(name)

